package com.ics.atable.chat.config;

import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

/**
 * Milvus配置类
 *
 * @author haohongbin
 */
@Slf4j
@Configuration
public class MilvusConfiguration {

    @Value("${spring.ai.vectorstore.milvus.client.host}")
    private String host;

    @Value("${spring.ai.vectorstore.milvus.client.port}")
    private String port;

    @Value("${spring.ai.vectorstore.milvus.client.username:}")
    private String username;

    @Value("${spring.ai.vectorstore.milvus.client.password:}")
    private String password;

    private final String dbName = "atable";

    @Bean
    public MilvusServiceClient milvusClient() {
        try {
            log.info("Connecting to Milvus at {}:{}", host, port);

            ConnectParam.Builder builder = ConnectParam.newBuilder()
                    .withUri("http://" + host + ":" + port)
                    .withConnectTimeout(5, TimeUnit.SECONDS)
                    .withKeepAliveTime(5, TimeUnit.SECONDS)
                    .withKeepAliveTimeout(5, TimeUnit.SECONDS)
                    .withIdleTimeout(5, TimeUnit.SECONDS);

            if (StringUtils.isNotBlank(username) && StringUtils.isNotBlank(password)) {
                builder.withAuthorization(username, password);
            }

            // 使用异步方式连接，设置超时
            CompletableFuture<MilvusServiceClient> future = CompletableFuture.supplyAsync(() -> {
                MilvusServiceClient client = new MilvusServiceClient(builder.build());

                // 使用更轻量级的方法测试连接
                try {
                    client.getVersion();
                    return client;
                } catch (Exception e) {
                    throw new RuntimeException("Connection test failed", e);
                }
            });

            // 等待连接结果，最多10秒
            MilvusServiceClient client = future.get(10, TimeUnit.SECONDS);
            log.info("Successfully connected to Milvus");
            return client;

        } catch (Exception e) {
            log.error("Failed to connect to Milvus at " + host + ":" + port, e);
            if (e instanceof java.util.concurrent.TimeoutException) {
                log.error("Connection to Milvus timed out after 10 seconds");
            }
            throw new RuntimeException("Failed to initialize Milvus client", e);
        }
    }

    @Bean
    public MilvusClientV2 milvusClientV2() {
        try {
            log.info("Connecting to Milvus V2 at {}:{}", host, port);

            // 使用异步方式连接，设置超时
            CompletableFuture<MilvusClientV2> future = CompletableFuture.supplyAsync(() -> {
                String uri = "http://" + host + ":" + port;
                return new MilvusClientV2(ConnectConfig.builder()
                        .uri(uri)
                        .dbName(dbName)
                        .username(username)
                        .password(password)
                        .connectTimeoutMs(5000)
                        .build());
            });

            // 等待连接结果，最多10秒
            MilvusClientV2 client = future.get(10, TimeUnit.SECONDS);
            log.info("Successfully connected to Milvus V2");
            return client;

        } catch (Exception e) {
            log.error("Failed to connect to Milvus V2 at " + host + ":" + port, e);
            if (e instanceof java.util.concurrent.TimeoutException) {
                log.error("Connection to Milvus V2 timed out after 10 seconds");
            }
            throw new RuntimeException("Failed to initialize Milvus V2 client", e);
        }
    }
}
